from enum import StrEnum
from typing import Literal

from langchain_community.chat_models import ChatTongyi
from langchain_core.prompts import ChatPromptTemplate
from pydantic import SecretStr, BaseModel, Field

chatLLM = ChatTongyi(
    model="qwen-plus-2025-04-28",
    # 此处以qwen-max为例，您可按需更换模型名称。模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
    streaming=True,
    api_key=SecretStr("sk-d16b46d66abb45bb960bd9c57804e2f9"),
    # other params...
)

# class Classification(BaseModel):
#     sentiment: str = Field(description="文本的情感")
#     aggressiveness: int = Field(
#         description="文本的攻击性程度从1到10，其中1表示非常低，10表示非常高"
#     )
#     language: str = Field(description="文本所使用的语言")


class Classification(BaseModel):
    sentiment: Literal["快乐", "平常心", "忧伤","愤怒"] = Field(description="文本的情感")
    aggressiveness: Literal[2,4,6,8] = Field(
        description="文本的攻击性程度,数字越低攻击程度越高",
    )
    language: Literal["中文", "英语", "日语", "法语","无法识别"] = Field(description="文本所使用的语言")

structured_llm = chatLLM.with_structured_output(Classification)



tagging_prompt = ChatPromptTemplate.from_template(
    """
从以下段落中提取所需信息。
仅提取“分类”功能中提到的属性。
文章:
{input}
"""
)

# inp = "这个社会太邪恶了，到处充满着不公平！"
# inp = "这个社会好美好呀，到处充满着爱，人和人之间友好的相处。"
#德语
inp = "Diese Gesellschaft ist wunderschön, überall voller Liebe und Freundschaft zwischen den Menschen."
prompt = tagging_prompt.invoke({"input": inp})
response = structured_llm.invoke(prompt)

print(response)
print("=========")
print(response.model_dump())